package cn.lsh.spark.sql;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;

import java.util.Arrays;
import java.util.List;

public class UserDefinedUDTF extends GenericUDTF {

	@Override
	public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentException {
		List<? extends StructField> inputFields = argOIs.getAllStructFieldRefs();
		if (inputFields.size() != 1) {
			throw new UDFArgumentException("splitStruct() takes only one argument");
		}
		ObjectInspector inputInspector = inputFields.get(0).getFieldObjectInspector();
		if (inputFields instanceof StringObjectInspector) {
			List<String> fieldNames = Arrays.asList("xing", "ming");
			List<ObjectInspector> fieldOIs = Arrays.asList(inputInspector, inputInspector);
			return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames, fieldOIs);
		} else {
			throw new RuntimeException("UserDefinedUDTF只接受String类型入参");
		}
	}

	@Override
	public void process(Object[] args) throws HiveException {
		if (args != null && args.length > 0) {
			for (int i = 0; i < args.length; i++) {
				forward(args[i].toString().split(" "));
			}
		}
	}

	@Override
	public void close() throws HiveException {

	}
}
